package utils

import (
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"unicode"

	"github.com/yanyiwu/gojieba"
)

var (
	jieba       *gojieba.Jieba
	stopWords   map[string]struct{}
	initialized bool
)

// 初始化分词器和停用词表
func init() {
	initialize()
}

func initialize() {
	if !initialized {
		dictDir := path.Join(filepath.Dir(os.Args[0]), "dict")
		jiebaPath := path.Join(dictDir, "jieba.dict.utf8")
		hmmPath := path.Join(dictDir, "hmm_model.utf8")
		userPath := path.Join(dictDir, "user.dict.utf8")
		idfPath := path.Join(dictDir, "idf.utf8")
		stopPath := path.Join(dictDir, "stop_words.utf8")
		jieba = gojieba.NewJieba(jiebaPath, hmmPath, userPath, idfPath, stopPath)
		// jieba = gojieba.NewJieba()
		loadStopWords()
		initialized = true
	}
}

// 加载停用词表
func loadStopWords() {
	stopWords = make(map[string]struct{})
	// 常见中文停用词
	words := []string{"的", "了", "和", "是", "在", "我", "有", "你", "他", "这", "那", "就", "也", "要", "会", "可以", "与", "等"}
	for _, word := range words {
		stopWords[word] = struct{}{}
	}
}

// ExtractKeyWord 提取关键词入口函数
// algorithm: "tf" 或 "textrank"
func ExtractKeyWord(content string, topN int, algorithm string) []string {
	initialize()

	// 先清洗文本
	cleaned := cleanText(content)

	// 根据算法选择提取方式
	var keywords []string
	switch algorithm {
	case "textrank":
		keywords = extractKeywordsByTextRank(cleaned, topN)
	default: // 默认使用TF算法
		keywords = extractKeywordsByTF(cleaned, topN)
	}

	return keywords
}

// cleanText 文本清洗
func cleanText(text string) string {
	// 去除标点符号
	f := func(r rune) bool {
		return unicode.IsPunct(r) || unicode.IsSymbol(r)
	}
	fields := strings.FieldsFunc(text, f)
	cleaned := strings.Join(fields, " ")

	// 转为小写(英文)
	return strings.ToLower(cleaned)
}

// extractKeywordsByTF TF算法实现
func extractKeywordsByTF(text string, topN int) []string {
	// 1. 分词
	words := jieba.Cut(text, true)

	// 2. 停用词过滤
	filtered := filterStopWords(words)

	// 3. 统计词频
	freq := make(map[string]int)
	for _, word := range filtered {
		if len(word) > 1 { // 忽略单字
			freq[word]++
		}
	}

	// 4. 按词频排序取TopN
	return getTopWords(freq, topN)
}

// extractKeywordsByTextRank TextRank算法实现
func extractKeywordsByTextRank(text string, topN int) []string {
	// 1. 分词
	words := jieba.Cut(text, true)

	// 2. 过滤停用词和单字
	filtered := filterStopWords(words)
	var validWords []string
	for _, word := range filtered {
		if len(word) > 1 {
			validWords = append(validWords, word)
		}
	}

	// 3. 构建共现窗口(窗口大小为5)
	windowSize := 5
	graph := make(map[string]map[string]float64)

	for i := 0; i < len(validWords); i++ {
		for j := i + 1; j < min(i+windowSize, len(validWords)); j++ {
			word1 := validWords[i]
			word2 := validWords[j]

			if _, ok := graph[word1]; !ok {
				graph[word1] = make(map[string]float64)
			}
			if _, ok := graph[word2]; !ok {
				graph[word2] = make(map[string]float64)
			}

			graph[word1][word2] += 1.0
			graph[word2][word1] += 1.0
		}
	}

	// 4. 初始化权重
	scores := make(map[string]float64)
	for word := range graph {
		scores[word] = 1.0
	}

	// 5. 迭代计算TextRank(10次迭代)
	dampingFactor := 0.85
	iterations := 10

	for i := 0; i < iterations; i++ {
		newScores := make(map[string]float64)
		for word := range scores {
			sum := 0.0
			for neighbor, weight := range graph[word] {
				sum += weight / float64(len(graph[neighbor])) * scores[neighbor]
			}
			newScores[word] = (1 - dampingFactor) + dampingFactor*sum
		}
		scores = newScores
	}

	// 6. 获取TopN关键词
	return getTopByScore(scores, topN)
}

// 辅助函数
func filterStopWords(words []string) []string {
	var filtered []string
	for _, word := range words {
		if _, ok := stopWords[word]; !ok {
			filtered = append(filtered, word)
		}
	}
	return filtered
}

func getTopWords(freq map[string]int, topN int) []string {
	type wordFreq struct {
		word string
		freq int
	}

	var wfList []wordFreq
	for word, freq := range freq {
		wfList = append(wfList, wordFreq{word, freq})
	}

	sort.Slice(wfList, func(i, j int) bool {
		return wfList[i].freq > wfList[j].freq
	})

	var result []string
	for i := 0; i < topN && i < len(wfList); i++ {
		result = append(result, wfList[i].word)
	}

	return result
}

func getTopByScore(scores map[string]float64, topN int) []string {
	type kv struct {
		Key   string
		Value float64
	}

	var ss []kv
	for k, v := range scores {
		ss = append(ss, kv{k, v})
	}

	sort.Slice(ss, func(i, j int) bool {
		return ss[i].Value > ss[j].Value
	})

	var result []string
	for i := 0; i < topN && i < len(ss); i++ {
		result = append(result, ss[i].Key)
	}

	return result
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
